#include "Basis.h"
#include "Memory.h"
#include "Network.h"

MINIPORT_SEND_NET_BUFFER_LISTS_HANDLER OrgSendNetBufferLists; // original send function
RECEIVE_NET_BUFFER_LISTS_HANDLER OrgReceiveNetBuferLists; // original recv function
NDIS_HANDLE NdisProtocolHandle; // protocol handle
NDIS_HANDLE NdisMiniportHandle; // miniport handle

LIST_ENTRY					g_NetworkPacketList;
NDIS_SPIN_LOCK				g_NetworkPacketSpinLock;
NPAGED_LOOKASIDE_LIST		g_NetworkPacketLookaside;
HANDLE						g_NetworkPacketEventHandler;
PKEVENT						g_NetworkPacketEvent;

VOID InitializeNetwork()
{
	NdisInitializeListHead(&g_NetworkPacketList);
	NdisAllocateSpinLock(&g_NetworkPacketSpinLock);
	ExInitializeNPagedLookasideList(&g_NetworkPacketLookaside, NULL, NULL, 0, sizeof(NETWORK_PACKET_ITEM), 'ten', 0);

	UNICODE_STRING networkPacketEventName;
	RtlInitUnicodeString(&networkPacketEventName, NETWORK_PACKET_EVENT_NAME);
	g_NetworkPacketEvent = IoCreateNotificationEvent(&networkPacketEventName, &g_NetworkPacketEventHandler);
	KeClearEvent(g_NetworkPacketEvent);
}

VOID DisposeNetwork()
{
	FreeList(&g_NetworkPacketList, &g_NetworkPacketSpinLock.SpinLock, &g_NetworkPacketLookaside, FreeNetworkPacketItemCallback);
	ExDeleteNPagedLookasideList(&g_NetworkPacketLookaside);
	ZwClose(g_NetworkPacketEventHandler);
}

VOID FreeNetworkPacketItemCallback(PLIST_ENTRY pListEntry)
{
	PNETWORK_PACKET_ITEM pNetworkPacketItem = CONTAINING_RECORD(pListEntry, NETWORK_PACKET_ITEM, ListEntry);
	NdisFreeMemory(pNetworkPacketItem->Data.Content, 0, 0);
}

BOOLEAN ProcessIoCtlForNetwork(ULONG ioControlCode, PDEVICE_OBJECT deviceObject, PIRP irp, PIO_STACK_LOCATION irpStack)
{
	UNREFERENCED_PARAMETER(deviceObject);
	UNREFERENCED_PARAMETER(irpStack);

	switch (ioControlCode)
	{
	case IOCTL_NETWORK_GET_PROTOCOL_HEADER:
		ProcessGetProtocolHeader(irp);
		break;
	case IOCTL_NETWORK_NDIS_HOOK:
		ProcessNdisHook(irp);
		break;
	case IOCTL_NETWORK_GET_NETWORK_PACKET:
		ProcessGetNetworkPacket(irp, irpStack);
		break;
	default:
		return FALSE;
	}
	return TRUE;
}

VOID ProcessGetProtocolHeader(PIRP irp)
{
	// init characteristics
	NDIS_STRING	protoName = NDIS_STRING_CONST("NDISPROT");
	NDIS_PROTOCOL_DRIVER_CHARACTERISTICS   protocolChar;
	NdisZeroMemory(&protocolChar, sizeof(NDIS_PROTOCOL_DRIVER_CHARACTERISTICS));

	protocolChar.Header.Type = NDIS_OBJECT_TYPE_PROTOCOL_DRIVER_CHARACTERISTICS,
	protocolChar.Header.Size = sizeof(NDIS_PROTOCOL_DRIVER_CHARACTERISTICS);
	protocolChar.Header.Revision = NDIS_PROTOCOL_DRIVER_CHARACTERISTICS_REVISION_1;

	protocolChar.MajorNdisVersion = 6;
	protocolChar.MinorNdisVersion = 0;
	protocolChar.Name = protoName;

	protocolChar.SetOptionsHandler = NULL;
	protocolChar.OpenAdapterCompleteHandlerEx = NdisprotOpenAdapterComplete;
	protocolChar.CloseAdapterCompleteHandlerEx = NdisprotCloseAdapterComplete;
	protocolChar.SendNetBufferListsCompleteHandler = NdisprotSendComplete;
	protocolChar.OidRequestCompleteHandler = NdisprotRequestComplete;
	protocolChar.StatusHandlerEx = NdisprotStatus;
	protocolChar.UninstallHandler = NULL;
	protocolChar.ReceiveNetBufferListsHandler = NdisprotReceiveNetBufferLists;
	protocolChar.NetPnPEventHandler = NdisprotPnPEventHandler;
	protocolChar.BindAdapterHandlerEx = NdisprotBindAdapter;
	protocolChar.UnbindAdapterHandlerEx = NdisprotUnbindAdapter;

	// register
	NDIS_HANDLE ndisProtocolHandle;
	NTSTATUS status = NdisRegisterProtocolDriver(NULL, &protocolChar, &ndisProtocolHandle);
	if (status != NDIS_STATUS_SUCCESS)
	{
		KdPrint(("failed to register protocol driver, error 0x%x.", status));
		irp->IoStatus.Status = STATUS_UNSUCCESSFUL;
		return;
	}

	// save to output
	PNDIS_PROTOCOL_BLOCK pBlock = (PNDIS_PROTOCOL_BLOCK)ndisProtocolHandle;
	PULONG_PTR buffer = (PULONG_PTR)irp->AssociatedIrp.SystemBuffer;
	*buffer = pBlock->NextProtocol;
	irp->IoStatus.Information = sizeof(ULONG_PTR);
	irp->IoStatus.Status = STATUS_SUCCESS;

	// unregister
	NdisDeregisterProtocolDriver(ndisProtocolHandle);
}

VOID ProcessNdisHook(PIRP irp)
{
	PNDIS_HOOK_REQUEST pNdisHookRequest = (PNDIS_HOOK_REQUEST)irp->AssociatedIrp.SystemBuffer;
	if (pNdisHookRequest->IsInstall)
	{
		InterlockedExchange((PLONG)pNdisHookRequest->HookSendAddress, (LONG)HookSendNetBufferLists);
		InterlockedExchange((PLONG)pNdisHookRequest->HookRecvAddress, (LONG)HookReceiveNetBuferLists);
		NdisProtocolHandle = pNdisHookRequest->NdisProtocolHandle;
		NdisMiniportHandle = pNdisHookRequest->NdisMiniportHandle;
		OrgSendNetBufferLists = (MINIPORT_SEND_NET_BUFFER_LISTS_HANDLER)pNdisHookRequest->SendFuncAddress;
		OrgReceiveNetBuferLists = (RECEIVE_NET_BUFFER_LISTS_HANDLER)pNdisHookRequest->RecvFuncAddress;

		KdPrint(("hooked ndis"));
	}
	else
	{
		InterlockedExchange((PLONG)pNdisHookRequest->HookSendAddress, (LONG)pNdisHookRequest->SendFuncAddress);
		InterlockedExchange((PLONG)pNdisHookRequest->HookRecvAddress, (LONG)pNdisHookRequest->RecvFuncAddress);
		NdisProtocolHandle = NULL;
		NdisMiniportHandle = NULL;
		OrgSendNetBufferLists = NULL;
		OrgReceiveNetBuferLists = NULL;

		KdPrint(("unhooked ndis"));
	}
	irp->IoStatus.Status = STATUS_SUCCESS;
}

VOID ProcessGetNetworkPacket(PIRP irp, PIO_STACK_LOCATION irpStack)
{
	ULONG outputBufferLength = irpStack->Parameters.DeviceIoControl.OutputBufferLength;
	PNETWORK_PACKET_ITEM_DATA buffer = (PNETWORK_PACKET_ITEM_DATA)irp->AssociatedIrp.SystemBuffer;

	PLIST_ENTRY pListEntry = NdisInterlockedRemoveHeadList(&g_NetworkPacketList, &g_NetworkPacketSpinLock);
	if (pListEntry != NULL)
	{
		PNETWORK_PACKET_ITEM pNetworkPacketItem = CONTAINING_RECORD(pListEntry, NETWORK_PACKET_ITEM, ListEntry);
		buffer->Size = pNetworkPacketItem->Data.Size;
		if (pNetworkPacketItem->Data.Size <= 2048)
		{
			// copy data
			buffer->Type = pNetworkPacketItem->Data.Type;
			RtlCopyMemory(&buffer->Content, pNetworkPacketItem->Data.Content, pNetworkPacketItem->Data.Size);
		}
		else
		{
			// insuffisant buffer size
			buffer->Type = -2;
		}
		
		// free item
		NdisFreeMemory(pNetworkPacketItem->Data.Content, 0, 0);
		ExFreeToNPagedLookasideList(&g_NetworkPacketLookaside, pNetworkPacketItem);
	}
	else
	{
		// no more
		buffer->Type = -1;

		// reset event
		KeClearEvent(g_NetworkPacketEvent);
	}

	irp->IoStatus.Information = outputBufferLength;
	irp->IoStatus.Status = STATUS_SUCCESS;
}

VOID HookSendNetBufferLists(
	_In_  NDIS_HANDLE             MiniportAdapterContext,
	_In_  PNET_BUFFER_LIST        NetBufferLists,
	_In_  NDIS_PORT_NUMBER        PortNumber,
	_In_  ULONG                   SendFlags
	)
{
	KdPrint(("calling HookSendNetBufferLists"));

	PNET_BUFFER_LIST nextBufferList = NetBufferLists;
	while (nextBufferList)
	{
		PVOID buffer = NULL;
		ULONG bufferSize = 0;

		// get raw buffers data
		if (CopyNblToBuffer(NdisMiniportHandle, nextBufferList, &buffer, &bufferSize) == NDIS_STATUS_SUCCESS)
		{
			KdPrint(("send data %d", bufferSize));
			
			PNETWORK_PACKET_ITEM pNetworkPacketItem = ExAllocateFromNPagedLookasideList(&g_NetworkPacketLookaside);
			if (pNetworkPacketItem != NULL)
			{
				// collect data & put into queue
				pNetworkPacketItem->Data.Type = 0;
				pNetworkPacketItem->Data.Size = bufferSize;
				pNetworkPacketItem->Data.Content = buffer;
				NdisInterlockedInsertTailList(&g_NetworkPacketList, &pNetworkPacketItem->ListEntry, &g_NetworkPacketSpinLock);

				// signal event
				KeSetEvent(g_NetworkPacketEvent, 0, FALSE);
			}
			else
			{
				KdPrint(("failed to allocate memory for network packet item in send"));
			}
		}
		else
		{
			KdPrint(("failed to copy nbl to buffer in send"));
		}

		// next
		nextBufferList = NET_BUFFER_LIST_NEXT_NBL(nextBufferList);
	}

	// call original send
	if (OrgSendNetBufferLists != NULL) OrgSendNetBufferLists(MiniportAdapterContext, NetBufferLists, PortNumber, SendFlags);
}

VOID HookReceiveNetBuferLists(
	_In_  NDIS_HANDLE             ProtocolBindingContext,
	_In_  PNET_BUFFER_LIST        NetBufferLists,
	_In_  NDIS_PORT_NUMBER        PortNumber,
	_In_  ULONG                   NumberOfNetBufferLists,
	_In_  ULONG                   ReceiveFlags
	)
{
	KdPrint(("calling HookReceiveNetBuferLists"));

	PNET_BUFFER_LIST nextBufferList = NetBufferLists;
	while (nextBufferList)
	{
		PVOID buffer = NULL;
		ULONG bufferSize = 0;

		// get raw buffers data
		if (CopyNblToBuffer(NdisProtocolHandle, nextBufferList, &buffer, &bufferSize) == NDIS_STATUS_SUCCESS)
		{
			KdPrint(("recv data %d", bufferSize));
			
			PNETWORK_PACKET_ITEM pNetworkPacketItem = ExAllocateFromNPagedLookasideList(&g_NetworkPacketLookaside);
			if (pNetworkPacketItem != NULL)
			{
				// collect data & put into queue
				pNetworkPacketItem->Data.Type = 1;
				pNetworkPacketItem->Data.Size = bufferSize;
				pNetworkPacketItem->Data.Content = buffer;
				NdisInterlockedInsertTailList(&g_NetworkPacketList, &pNetworkPacketItem->ListEntry, &g_NetworkPacketSpinLock);

				// signal event
				KeSetEvent(g_NetworkPacketEvent, 0, FALSE);
			}
			else
			{
				KdPrint(("failed to allocate memory for network packet item in recv"));
			}
		}
		else
		{
			KdPrint(("failed to copy nbl to buffer in recv"));
		}

		// next
		nextBufferList = NET_BUFFER_LIST_NEXT_NBL(nextBufferList);
	}

	// call original recv
	if (OrgReceiveNetBuferLists != NULL) OrgReceiveNetBuferLists(ProtocolBindingContext, NetBufferLists, PortNumber, NumberOfNetBufferLists, ReceiveFlags);
}

//////////////////////////////////////////////////////////////////////////
// Dummy Handlers
//////////////////////////////////////////////////////////////////////////

VOID
NdisprotOpenAdapterComplete(
IN NDIS_HANDLE                  ProtocolBindingContext,
IN NDIS_STATUS                  Status
)
/*++

Routine Description:

Completion routine called by NDIS if our call to NdisOpenAdapterEx
pends. Wake up the thread that called NdisOpenAdapterEx.

Arguments:

ProtocolBindingContext - pointer to open context structure
Status - status of the open

Return Value:

None

--*/
{
	UNREFERENCED_PARAMETER(ProtocolBindingContext);
	UNREFERENCED_PARAMETER(Status);
}

VOID
NdisprotCloseAdapterComplete(
IN NDIS_HANDLE                  ProtocolBindingContext
)
/*++

Routine Description:

Called by NDIS to complete a pended call to NdisCloseAdapter.
We wake up the thread waiting for this completion.

Arguments:

ProtocolBindingContext - pointer to open context structure

Return Value:

None

--*/
{
	UNREFERENCED_PARAMETER(ProtocolBindingContext);
}

VOID
NdisprotSendComplete(
IN NDIS_HANDLE                  ProtocolBindingContext,
IN PNET_BUFFER_LIST             pNetBufferList,
IN ULONG                        SendCompleteFlags
)
/*++

Routine Description:

NDIS entry point called to signify completion of a packet send.
We pick up and complete the Write IRP corresponding to this packet.

Arguments:

ProtocolBindingContext - pointer to open context
pNetBufferList - NetBufferList that completed send
SendCompleteFlags - Specifies if the caller is at DISPATCH level

Return Value:

None

--*/
{
	UNREFERENCED_PARAMETER(ProtocolBindingContext);
	UNREFERENCED_PARAMETER(pNetBufferList);
	UNREFERENCED_PARAMETER(SendCompleteFlags);
}

VOID
NdisprotRequestComplete(
IN NDIS_HANDLE                  ProtocolBindingContext,
IN PNDIS_OID_REQUEST            pNdisRequest,
IN NDIS_STATUS                  Status
)
/*++

Routine Description:

NDIS entry point indicating completion of a pended NDIS_REQUEST.

Arguments:

ProtocolBindingContext - pointer to open context
pNdisRequest - pointer to NDIS request
Status - status of reset completion

Return Value:

None

--*/
{
	UNREFERENCED_PARAMETER(ProtocolBindingContext);
	UNREFERENCED_PARAMETER(pNdisRequest);
	UNREFERENCED_PARAMETER(Status);
}

VOID
NdisprotStatus(
IN NDIS_HANDLE                  ProtocolBindingContext,
IN PNDIS_STATUS_INDICATION      StatusIndication
)
/*++

Routine Description:

Protocol entry point called by NDIS to indicate a change
in status at the miniport.

We make note of reset and media connect status indications.

Arguments:

ProtocolBindingContext - pointer to open context
StatusIndication - pointer to NDIS_STATUS_INDICATION

Return Value:

None

--*/
{
	UNREFERENCED_PARAMETER(ProtocolBindingContext);
	UNREFERENCED_PARAMETER(StatusIndication);
}

VOID
NdisprotReceiveNetBufferLists(
IN NDIS_HANDLE                  ProtocolBindingContext,
IN PNET_BUFFER_LIST             pNetBufferLists,
IN NDIS_PORT_NUMBER             PortNumber,
IN ULONG                        NumberOfNetBufferLists,
IN ULONG                        ReceiveFlags
)
/*++

Routine Description:

Protocol entry point called by NDIS if the driver below
uses NDIS 6 net buffer list indications.

If the miniport allows us to hold on to this net buffer list, we
use it as is, otherwise we make a copy.

Arguments:

ProtocolBindingContext - pointer to open context
pNetBufferLists - a list of the Net Buffer lists being indicated up.
PortNumber - Port on which the Net Bufer list was received
NumberOfNetBufferLists - the number of NetBufferLists in this indication
ReceiveFlags - indicates whether the NetBufferLists can be pended in
the protocol driver.

Return Value:

--*/
{
	UNREFERENCED_PARAMETER(ProtocolBindingContext);
	UNREFERENCED_PARAMETER(pNetBufferLists);
	UNREFERENCED_PARAMETER(PortNumber);
	UNREFERENCED_PARAMETER(NumberOfNetBufferLists);
	UNREFERENCED_PARAMETER(ReceiveFlags);
}

NDIS_STATUS
NdisprotPnPEventHandler(
IN NDIS_HANDLE                  ProtocolBindingContext,
IN PNET_PNP_EVENT_NOTIFICATION  pNetPnPEventNotification
)
/*++

Routine Description:

Called by NDIS to notify us of a PNP event. The most significant
one for us is power state change.

Arguments:

ProtocolBindingContext - pointer to open context structure
this is NULL for global reconfig events.

pNetPnPEventNotification - pointer to the PNP event notification

Return Value:

Our processing status for the PNP event.

--*/
{
	UNREFERENCED_PARAMETER(ProtocolBindingContext);
	UNREFERENCED_PARAMETER(pNetPnPEventNotification);
	return NDIS_STATUS_SUCCESS;
}

NDIS_STATUS
NdisprotBindAdapter(
IN NDIS_HANDLE                  ProtocolDriverContext,
IN NDIS_HANDLE                  BindContext,
IN PNDIS_BIND_PARAMETERS        BindParameters
)
/*++

Routine Description:

Protocol Bind Handler entry point called when NDIS wants us
to bind to an adapter. We go ahead and set up a binding.
An OPEN_CONTEXT structure is allocated to keep state about
this binding.

Arguments:


Return Value:

None

--*/
{
	UNREFERENCED_PARAMETER(ProtocolDriverContext);
	UNREFERENCED_PARAMETER(BindContext);
	UNREFERENCED_PARAMETER(BindParameters);
	return NDIS_STATUS_SUCCESS;
}

NDIS_STATUS
NdisprotUnbindAdapter(
IN NDIS_HANDLE                  UnbindContext,
IN NDIS_HANDLE                  ProtocolBindingContext
)
/*++

Routine Description:

NDIS calls this when it wants us to close the binding to an adapter.

Arguments:

ProtocolBindingContext - pointer to open context structure
UnbindContext - to use in NdisCompleteUnbindAdapter if we return pending

Return Value:

pending or success

--*/
{
	UNREFERENCED_PARAMETER(UnbindContext);
	UNREFERENCED_PARAMETER(ProtocolBindingContext);
	return NDIS_STATUS_SUCCESS;
}
